package com.cwj.algorithm.graph;

import com.cwj.algorithm.priority.IndexPriorityQueue;

import java.util.Arrays;
import java.util.Deque;
import java.util.LinkedList;

/**
 * @author chenwujie
 * @date 2020-12-24 13:29
 */
public class PrimMST {
    private EdgeWeightedGraph edgeWeightedGraph;
    /**
     * 到达顶点的最短边
     */
    private Edge[] edgeTo;
    /**
     * 到达顶点的最短边的权重
     */
    private int[] distTo;
    /**
     * 标识顶点是否已经存在最小生成树中
     */
    private boolean[] marked;
    /**
     * 使用最小索引优先队列维护有效横切边的权重值
     */
    private IndexPriorityQueue<Integer> pq;

    public PrimMST(EdgeWeightedGraph edgeWeightedGraph){
        this.edgeWeightedGraph = edgeWeightedGraph;
        this.edgeTo = new Edge[edgeWeightedGraph.getV()];
        this.distTo = new int[edgeWeightedGraph.getV()];
        Arrays.fill(distTo, Integer.MAX_VALUE);
        this.marked = new boolean[edgeWeightedGraph.getV()];
        this.pq = new IndexPriorityQueue<>(edgeWeightedGraph.getV(), true);

    }

    /**
     * 生成树
     */
    public void mst(){
        // 默认顶点0 为第一个顶点
        distTo[0] = 0;
        pq.insert(0, 0);
        while(!pq.isEmpty()){
            visit(pq.removeFirst());
        }
    }

    /**
     * 获取最小生成树的所有边
     * @return
     */
    public Deque<Edge> edges(){
        Deque<Edge> edges = new LinkedList<>();
        for (Edge edge : edgeTo) {
            edges.push(edge);
        }
        return edges;
    }

    /**
     * 添加顶点v到最小生成树中
     * @param v
     */
    private void visit(int v){
        // 标识顶点v已经在最小生成树中
        marked[v] = true;
        // 获取顶点v的邻接表
        Deque<Edge> edges = edgeWeightedGraph.edgesByV(v);
        // 遍历邻接表
        for (Edge e : edges) {
            // 获取边的另一顶点
            int w = e.other(v);

            if(marked[w]){
                // 如果该顶点已经存在在最小生成树中，则不处理
                continue;
            }

            if(e.getWeight() < distTo[w]){
                // 更新 到w顶点的最小边
                edgeTo[w] = e;
                distTo[w] = e.getWeight();

                if(pq.contains(w)){
                    // 更新到w顶点的横切边的值
                    pq.change(w, e.getWeight());
                }else{
                    // 添加 到w顶点的横切边
                    pq.insert(w, e.getWeight());
                }
            }

        }

    }

}
